#include "SkyNet.h"

static layer config[layer_count] = {
{ "conv0", 320,160,32, 320,160,32, 0,0,0 },  //conv0
{ "conv1", 320,160,32, 320,160,32, 3,1,1 },  //conv1
{ "conv2", 320,160,32, 320,160,64, 1,1,0 },  //conv2
{ "pool1", 320,160,64, 160,80,64,  2,2,0 },  //pool1
{ "conv3", 160,80,64,  160,80,64,  3,1,1 },  //conv3
{ "conv4", 160,80,64,  160,80,96,  1,1,0 },  //conv4
{ "pool2", 160,80,96,  80,40,96,   2,2,0 },  //pool2
{ "conv5", 80,40,96,   80,40,96,   3,1,1 },  //conv5
{ "conv6", 80,40,96,   80,40,192,  1,1,0 },  //conv6
{ "reorg", 80,40,192,  40,20,768,  2,2,0 },  //reorg
{ "pool3", 80,40,192,  40,20,192,  2,2,0 },  //pool3
{ "conv7", 40,20,192,  40,20,192,  3,1,1 },  //conv7
{ "conv8", 40,20,192,  40,20,384,  1,1,0 },  //conv8
{ "conv9", 40,20,384,  40,20,384,  3,1,1 },  //conv9
{ "conv10",40,20,384,  40,20,512,  1,1,0 },  //conv10
{ "cat",   40,20,192,  40,20,1280, 0,0,0 },  //concat
{ "conv11",40,20,1280, 40,20,1280, 3,1,1 },  //conv11
{ "conv12",40,20,1280, 40,20,96,   1,1,0 },  //conv12
{ "conv13",40,20,96,   40,20,32,   1,1,0 },  //conv13
};

void write_result_to_ofm(ADT32 *ofm)
{
    const int input_time = 11;
    int offset[input_time] = {pool1_offset, pool2_offset, conv5_offset, conv6_offset, pool3_offset, conv7_offset, conv8_offset, conv9_offset, conv10_offset,
    conv11_offset, conv12_offset};
    int config_index[] = {3, 6, 7, 8, 10, 11, 12, 13, 14, 16, 17};
    for(int i = 0; i< input_time; i++)
    {
        layer l = config[config_index[i]];
        int len = l.oc*l.ow*l.oh;
        ADT *tmp[4];
        for (int p = 0; p < 4; p++)
        {
            tmp[p] = (ADT *)sds_alloc(sizeof(ADT)*len);
            load_fm(tmp[p], l);
        }

        ADT *tmp_blob;
        tmp_blob = (ADT *)sds_alloc(64 * 643 * 323 * sizeof(ADT));

        for (int i = 0; i < 64 ; i++)
        {
            for (int j = 0; j < 323; j++)
            {
                for (int k = 0; k < 643; k++)
                {
                    tmp_blob[i*323*643+j*643+k] = -128;
                }
            }
        }

        stitch(tmp, tmp_blob, l);
        fm_DT_2_DT32(tmp_blob, &ofm[offset[i]], l);

		for (int p = 0; p < 4; p++)
		{
			sds_free(tmp[p]);
		}
		sds_free(tmp_blob);
    }
}

int main() {

//*************************************init *********************************
	std::cout << "init SkyNet" << std::endl;
	WDT32* weight;
	BDT8* biasm;
	BDT8* bbox_blob8;
	BDT* bbox_blob;
	ADT* data[4];
	ADT* data_blob;
	ADT32* data_blob32;
	ADT* ofm_blob;
	ADT32* ofm_blob32;
	ADT* ofm[4];
	BDT* bbox[4];

	for(int p=0; p<4; p++)
	{
		data[p] = (ADT*)sds_alloc(32*160*320*sizeof(ADT));
		ofm[p] = (ADT*)sds_alloc(64*320*640*sizeof(ADT));
		bbox[p] = (BDT*)sds_alloc(20*40*16*sizeof(BDT));
	}
	data_blob = (ADT*)sds_alloc(32*323*643*sizeof(ADT));
	data_blob32 = (ADT32*)sds_alloc(32*323*643*sizeof(ADT));
	weight = (WDT32*)sds_alloc(443344*sizeof(WDT));
	biasm = (BDT8*)sds_alloc(6848*sizeof(BDT));
	bbox_blob8 = (BDT8*)sds_alloc(43*83*32*sizeof(BDT));
	bbox_blob = (BDT*)sds_alloc(43*83*32*sizeof(BDT));
	ofm_blob32 = (ADT32*)sds_alloc(32*conv13_offset*sizeof(ADT));
	ofm_blob = (ADT*)sds_alloc(64*643*323*sizeof(ADT));

	if (data_blob ==NULL || data_blob32 ==NULL || weight ==NULL || biasm ==NULL || 
	bbox_blob8 ==NULL || bbox_blob ==NULL || ofm_blob32 ==NULL || ofm_blob ==NULL)
	{
		std::cout << "ERROR allocate memory! " << std::endl;
		return -1;
	}
	else
	{
		std::cout << "Success allocate memory! " << std::endl;
	}	

	for(int i=0; i<conv13_offset; i++)
	{
		for(int j=0; j<32; j++)
		{
		#ifndef DEBUG
			ofm_blob32[i].range(8*j+7,8*j) = -128;
		#else 
			ofm_blob32[i].data[j] = -128;
		#endif
		}
	}

	//*************************************load data *********************************
	std::cout << "load parameter" << std::endl;
	load_weight(weight, 443344);
	load_biasm(biasm, 6848);

	std::cout << "load image" << std::endl;
	for (int p = 0; p<4; p++)
		load_fm(data[p], config[0]);

	std::cout << "stitch image" << std::endl;
	stitch(data, data_blob, config[0]);

	std::cout << "transform image" << std::endl;
	fm_DT_2_DT32(data_blob, data_blob32, config[0]);

	//*************************************HLS, Skynet *********************************
	std::cout << "load correct result" << std::endl;
	// write_result_to_ofm(ofm_blob32);
	
	std::cout << "SkyNet start" << std::endl;
	SkyNet(data_blob32, ofm_blob32, weight, biasm, bbox_blob8);

	//************************************* check the results*********************************
	std::cout << "check SkyNet results" << std::endl;

	fm_DT32_2_DT(&ofm_blob32[pool1_offset], ofm_blob, config[3]);
	distitch(ofm_blob, ofm, config[3]);
	for (int p = 0; p<4; p++)
	{
		check_fm(ofm[p], config[3]);
	}
	std::cout << std::endl;

	fm_DT32_2_DT(&ofm_blob32[pool2_offset], ofm_blob, config[6]);
	distitch(ofm_blob, ofm, config[6]);
	for (int p = 0; p<4; p++)
	{
		check_fm(ofm[p], config[6]);
	}
	std::cout << std::endl;

	fm_DT32_2_DT(&ofm_blob32[conv5_offset], ofm_blob, config[7]);
	distitch(ofm_blob, ofm, config[7]);
	for (int p = 0; p<4; p++)
	{
		check_fm(ofm[p], config[7]);
	}
	std::cout << std::endl;

	fm_DT32_2_DT(&ofm_blob32[conv6_offset], ofm_blob, config[8]);
	distitch(ofm_blob, ofm, config[8]);
	for (int p = 0; p<4; p++)
	{
		check_fm(ofm[p], config[8]);
	}
	std::cout << std::endl;

	fm_DT32_2_DT(&ofm_blob32[pool3_offset], ofm_blob, config[10]);
	distitch(ofm_blob, ofm, config[10]);
	for (int p = 0; p<4; p++)
	{
		check_fm(ofm[p], config[10]);
	}
	std::cout << std::endl;

	fm_DT32_2_DT(&ofm_blob32[conv7_offset], ofm_blob, config[11]);
	distitch(ofm_blob, ofm, config[11]);
	for (int p = 0; p<4; p++)
	{
		check_fm(ofm[p], config[11]);
	}
	std::cout << std::endl;

	fm_DT32_2_DT(&ofm_blob32[conv8_offset], ofm_blob, config[12]);
	distitch(ofm_blob, ofm, config[12]);
	for (int p = 0; p<4; p++)
	{
		check_fm(ofm[p], config[12]);
	}
	std::cout << std::endl;

	fm_DT32_2_DT(&ofm_blob32[conv9_offset], ofm_blob, config[13]);
	distitch(ofm_blob, ofm, config[13]);
	for (int p = 0; p<4; p++)
	{
		check_fm(ofm[p], config[13]);
	}
	std::cout << std::endl;

	fm_DT32_2_DT(&ofm_blob32[conv10_offset], ofm_blob, config[14]);
	distitch(ofm_blob, ofm, config[14]);
	for (int p = 0; p<4; p++)
	{
		check_fm(ofm[p], config[14]);
	}
	std::cout << std::endl;

	fm_DT32_2_DT(&ofm_blob32[conv11_offset], ofm_blob, config[16]);
	distitch(ofm_blob, ofm, config[16]);
	for (int p = 0; p<4; p++)
	{
		check_fm(ofm[p], config[16]);
	}
	std::cout << std::endl;

	fm_DT32_2_DT(&ofm_blob32[conv12_offset], ofm_blob, config[17]);
	distitch(ofm_blob, ofm, config[17]);
	for (int p = 0; p<4; p++)
	{
		check_fm(ofm[p], config[17]);
	}
	std::cout << std::endl;

	bbox_DT8_2_DT(bbox_blob8, bbox_blob, config[18]);
	distitch_bbox(bbox_blob, bbox, config[18]);
	for (int p = 0; p<4; p++)
	{
		check_bbox(bbox[p], config[18]);
	}
	std::cout << std::endl;

	for (int p = 0; p<4; p++)
    {
        sds_free(data[p]);
        sds_free(ofm[p]);
        sds_free(bbox[p]);
    }
    sds_free(data_blob);
    sds_free(data_blob32);
    sds_free(weight);
    sds_free(biasm);
    sds_free(bbox_blob8);
    sds_free(bbox_blob);
    sds_free(ofm_blob32);
    sds_free(ofm_blob);
}
